import cv2
import numpy as np
import os
import logging
from pathlib import Path
import argparse
from tqdm import tqdm

from edge_detection import apply_canny_edge_detection
from kmeans_clustering import apply_kmeans_clustering

# 设置日志
def setup_logger():
    logger = logging.getLogger('THWImageProcessor')
    if not logger.handlers:
        logger.setLevel(logging.INFO)
        # 创建控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        logger.addHandler(console_handler)
        
        # 创建文件处理器
        file_handler = logging.FileHandler('thw_processing.log')
        file_handler.setLevel(logging.INFO)
        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        logger.addHandler(file_handler)
    return logger

logger = setup_logger()

def optimize_results(image, edge_image, clustered_image, layer_masks):
    """
    优化聚类结果和边缘检测结果
    
    Args:
        image: 原始图像
        edge_image: 边缘检测结果（白色背景黑色边缘）
        clustered_image: 聚类结果图像
        layer_masks: 聚类层掩码列表
        
    Returns:
        优化后的结果图像
    """
    logger.info("开始优化聚类结果和边缘...")
    
    # 1. 将边缘图像转换为掩码（黑色=边缘，白色=非边缘）
    edge_mask = edge_image < 128
    
    # 2. 创建优化后的聚类结果图像
    optimized_image = clustered_image.copy()
    
    # 3. 将边缘区域标记为黑色
    optimized_image[edge_mask] = [0, 0, 0]
    
    # 4. 对每个聚类层进行优化
    optimized_layers = []
    optimized_masks = []
    
    for i, mask in enumerate(layer_masks):
        # 4.1 创建新掩码，排除边缘区域
        optimized_mask = np.logical_and(mask > 0, ~edge_mask).astype(np.uint8) * 255
        
        # 4.2 对掩码进行形态学处理，去除噪点
        kernel = np.ones((3,3), np.uint8)
        optimized_mask = cv2.morphologyEx(optimized_mask, cv2.MORPH_OPEN, kernel)
        optimized_mask = cv2.morphologyEx(optimized_mask, cv2.MORPH_CLOSE, kernel)
        
        # 4.3 处理小连通区域
        num_labels, labels = cv2.connectedComponents(optimized_mask)
        for label in range(1, num_labels):
            area = np.sum(labels == label)
            if area < 100:  # 小于100像素的区域视为噪点
                optimized_mask[labels == label] = 0
        
        # 4.4 创建优化后的图层
        optimized_layer = np.full_like(image, 255)
        layer_pixels = optimized_mask > 0
        
        if np.any(layer_pixels):
            # 提取该区域的原始颜色
            region_colors = image[layer_pixels]
            # 中值滤波平滑颜色
            smoothed_colors = cv2.medianBlur(region_colors.reshape(-1, 3), 3).reshape(region_colors.shape)
            # 应用到图层
            optimized_layer[layer_pixels] = smoothed_colors
        
        optimized_layers.append(optimized_layer)
        optimized_masks.append(optimized_mask)
    
    logger.info(f"优化完成，共 {len(optimized_layers)} 个优化后的聚类层")
    return optimized_image, optimized_layers, optimized_masks

def process_single_image(image_path, output_dir, low_threshold=30, high_threshold=90):
    """
    处理单张图像
    
    Args:
        image_path: 图像路径
        output_dir: 输出目录
        low_threshold: Canny低阈值
        high_threshold: Canny高阈值
        
    Returns:
        成功返回True，失败返回False
    """
    try:
        # 读取图像
        logger.info(f"读取图像: {image_path}")
        image = cv2.imread(str(image_path))
        if image is None:
            logger.error(f"无法读取图像: {image_path}")
            return False
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"输出目录: {output_dir}")
        
        # 保存原始图像
        original_output_path = os.path.join(output_dir, "1_原图.png")
        cv2.imwrite(original_output_path, image)
        
        # 执行Canny边缘检测，使用指定的阈值
        logger.info(f"执行Canny边缘检测 (低阈值={low_threshold}, 高阈值={high_threshold})...")
        edges = apply_canny_edge_detection(image, low_threshold=low_threshold, high_threshold=high_threshold)
        edge_output_path = os.path.join(output_dir, "2_边缘检测.png")
        cv2.imwrite(edge_output_path, edges)
        
        # 反转边缘图像用于KMeans处理（黑色背景白色边缘）
        edges_for_kmeans = cv2.bitwise_not(edges)
        
        # 执行K-means++聚类
        logger.info("执行K-means++聚类...")
        result = apply_kmeans_clustering(image)
        if result is None:
            logger.error("聚类处理失败")
            return False
            
        clustered_image, layer_results, layer_masks = result
        
        # 保存聚类结果
        clustered_output_path = os.path.join(output_dir, "3_聚类结果.png")
        cv2.imwrite(clustered_output_path, clustered_image)
        
        # 保存每个图层和掩码
        logger.info(f"保存 {len(layer_results)} 个聚类图层...")
        for i, (layer, mask) in enumerate(zip(layer_results, layer_masks)):
            layer_output_path = os.path.join(output_dir, f"4_{i+1}_图层.png")
            mask_output_path = os.path.join(output_dir, f"5_{i+1}_掩码.png")
            cv2.imwrite(layer_output_path, layer)
            cv2.imwrite(mask_output_path, mask)
        
        # 优化聚类结果和边缘
        logger.info("优化聚类结果和边缘...")
        optimized_image, optimized_layers, optimized_masks = optimize_results(
            image, edges, clustered_image, layer_masks
        )
        
        # 保存优化后的结果
        optimized_output_path = os.path.join(output_dir, "6_优化结果.png")
        cv2.imwrite(optimized_output_path, optimized_image)
        
        # 保存优化后的图层和掩码
        for i, (layer, mask) in enumerate(zip(optimized_layers, optimized_masks)):
            layer_output_path = os.path.join(output_dir, f"7_{i+1}_优化图层.png")
            mask_output_path = os.path.join(output_dir, f"8_{i+1}_优化掩码.png")
            cv2.imwrite(layer_output_path, layer)
            cv2.imwrite(mask_output_path, mask)
        
        logger.info("图像处理完成")
        return True
        
    except Exception as e:
        logger.error(f"处理图像时发生错误: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False

def batch_process_directory(input_folder, output_folder, low_threshold=30, high_threshold=90):
    """
    批量处理文件夹中的所有图像
    
    Args:
        input_folder: 输入文件夹路径
        output_folder: 输出文件夹路径
        low_threshold: Canny低阈值
        high_threshold: Canny高阈值
        
    Returns:
        成功处理的图像数量
    """
    input_folder = Path(input_folder)
    output_folder = Path(output_folder)
    
    if not input_folder.exists() or not input_folder.is_dir():
        logger.error(f"输入文件夹不存在或不是有效目录: {input_folder}")
        return 0
    
    # 创建输出目录
    os.makedirs(output_folder, exist_ok=True)
    
    # 获取所有图像文件
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']
    image_files = []
    for ext in image_extensions:
        image_files.extend(list(input_folder.glob(f"*{ext}")))
        image_files.extend(list(input_folder.glob(f"*{ext.upper()}")))
    
    if not image_files:
        logger.warning(f"输入文件夹 {input_folder} 中没有找到图像文件")
        return 0
    
    logger.info(f"找到 {len(image_files)} 个图像文件，开始批量处理...")
    
    success_count = 0
    # 使用tqdm显示进度条
    for image_path in tqdm(image_files, desc="处理图像"):
        # 为每张图像创建一个单独的输出目录
        image_output_dir = output_folder / image_path.stem
        if process_single_image(
            image_path, 
            image_output_dir, 
            low_threshold=low_threshold, 
            high_threshold=high_threshold
        ):
            success_count += 1
    
    logger.info(f"批量处理完成，成功处理 {success_count}/{len(image_files)} 个图像")
    return success_count

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="桃花坞论文图像处理工具")
    parser.add_argument("--input", "-i", type=str, default=r"D:\桃花坞论文", help="输入图像或文件夹路径")
    parser.add_argument("--output", "-o", type=str, default=None, help="输出目录路径")
    parser.add_argument("--low", "-l", type=int, default=30, help="Canny边缘检测低阈值")
    parser.add_argument("--high", "-h", type=int, default=90, help="Canny边缘检测高阈值")
    parser.add_argument("--single", "-s", type=str, default=None, help="处理单张图像")
    
    args = parser.parse_args()
    
    # 设置输出目录
    if args.output is None:
        args.output = Path(args.input) / "output"
    
    # 处理单张图像
    if args.single:
        single_image_path = Path(args.input) / args.single
        if not single_image_path.exists():
            logger.error(f"指定的图像不存在: {single_image_path}")
        else:
            output_dir = Path(args.output) / single_image_path.stem
            process_single_image(
                single_image_path, 
                output_dir, 
                low_threshold=args.low, 
                high_threshold=args.high
            )
    # 批量处理
    else:
        batch_process_directory(
            args.input, 
            args.output, 
            low_threshold=args.low, 
            high_threshold=args.high
        ) 